"""
音频质量增强算法
提供各种音频增强和优化功能
"""

import numpy as np
import librosa
import scipy.signal as signal
from typing import Dict, Any, Optional, Tuple
from dataclasses import dataclass
from enum import Enum
import io
import tempfile
import soundfile as sf


class EnhancementType(Enum):
    """增强类型"""
    NOISE_REDUCTION = "noise_reduction"
    EQUALIZATION = "equalization"
    DYNAMICS = "dynamics"
    STEREO_ENHANCEMENT = "stereo_enhancement"
    PITCH_CORRECTION = "pitch_correction"
    TEMPO_ADJUSTMENT = "tempo_adjustment"
    REVERBERATION = "reverberation"
    CLARITY = "clarity"


@dataclass
class EnhancementSettings:
    """增强设置"""
    noise_reduction: bool = False
    noise_reduction_strength: float = 0.5
    equalization: bool = False
    eq_bands: Dict[str, float] = None
    dynamics_compression: bool = False
    compression_ratio: float = 4.0
    compression_threshold: float = -20.0
    stereo_enhancement: bool = False
    stereo_width: float = 1.2
    pitch_correction: bool = False
    pitch_shift: float = 0.0
    tempo_adjustment: bool = False
    tempo_factor: float = 1.0
    reverb: bool = False
    reverb_room_size: float = 0.3
    reverb_damping: float = 0.5
    clarity_enhancement: bool = False
    clarity_strength: float = 0.5


@dataclass
class EnhancementResult:
    """增强结果"""
    success: bool
    enhanced_data: Optional[np.ndarray] = None
    sample_rate: int = 0
    metadata: Optional[Dict[str, Any]] = None
    error_message: Optional[str] = None


class AudioEnhancer:
    """音频增强器"""
    
    def __init__(self):
        self.default_eq_bands = {
            'low': 1.0,      # 60-250 Hz
            'low_mid': 1.0,  # 250-500 Hz
            'mid': 1.0,      # 500-2000 Hz
            'high_mid': 1.0, # 2000-4000 Hz
            'high': 1.0      # 4000-16000 Hz
        }
    
    def enhance_audio(
        self,
        data: np.ndarray,
        sample_rate: int,
        settings: EnhancementSettings
    ) -> EnhancementResult:
        """音频增强处理"""
        try:
            enhanced_data = data.copy()
            metadata = {
                'original_sample_rate': sample_rate,
                'enhancements_applied': []
            }
            
            # 应用各种增强效果
            if settings.noise_reduction:
                enhanced_data = self._apply_noise_reduction(
                    enhanced_data, sample_rate, settings.noise_reduction_strength
                )
                metadata['enhancements_applied'].append('noise_reduction')
            
            if settings.equalization:
                eq_bands = settings.eq_bands or self.default_eq_bands
                enhanced_data = self._apply_equalization(
                    enhanced_data, sample_rate, eq_bands
                )
                metadata['enhancements_applied'].append('equalization')
                metadata['eq_settings'] = eq_bands
            
            if settings.dynamics_compression:
                enhanced_data = self._apply_dynamics_compression(
                    enhanced_data, settings.compression_ratio, settings.compression_threshold
                )
                metadata['enhancements_applied'].append('dynamics_compression')
            
            if settings.stereo_enhancement:
                enhanced_data = self._apply_stereo_enhancement(
                    enhanced_data, settings.stereo_width
                )
                metadata['enhancements_applied'].append('stereo_enhancement')
            
            if settings.pitch_correction:
                enhanced_data = self._apply_pitch_correction(
                    enhanced_data, sample_rate, settings.pitch_shift
                )
                metadata['enhancements_applied'].append('pitch_correction')
            
            if settings.tempo_adjustment:
                enhanced_data = self._apply_tempo_adjustment(
                    enhanced_data, sample_rate, settings.tempo_factor
                )
                metadata['enhancements_applied'].append('tempo_adjustment')
            
            if settings.reverb:
                enhanced_data = self._apply_reverberation(
                    enhanced_data, sample_rate, settings.reverb_room_size, settings.reverb_damping
                )
                metadata['enhancements_applied'].append('reverberation')
            
            if settings.clarity_enhancement:
                enhanced_data = self._apply_clarity_enhancement(
                    enhanced_data, sample_rate, settings.clarity_strength
                )
                metadata['enhancements_applied'].append('clarity_enhancement')
            
            return EnhancementResult(
                success=True,
                enhanced_data=enhanced_data,
                sample_rate=sample_rate,
                metadata=metadata
            )
            
        except Exception as e:
            return EnhancementResult(
                success=False,
                error_message=f"Enhancement failed: {str(e)}"
            )
    
    def _apply_noise_reduction(
        self, data: np.ndarray, sample_rate: int, strength: float
    ) -> np.ndarray:
        """噪声抑制"""
        # 使用谱减法进行噪声抑制
        stft = librosa.stft(data)
        magnitude = np.abs(stft)
        phase = np.angle(stft)
        
        # 估计噪声谱（使用前几帧）
        noise_frames = int(0.1 * sample_rate / 512)  # 前0.1秒
        noise_magnitude = np.mean(magnitude[:, :noise_frames], axis=1, keepdims=True)
        
        # 谱减法
        alpha = 2.0 * strength  # 过减因子
        magnitude_denoised = magnitude - alpha * noise_magnitude
        magnitude_denoised = np.maximum(magnitude_denoised, 0.1 * magnitude)
        
        # 重构信号
        stft_denoised = magnitude_denoised * np.exp(1j * phase)
        data_denoised = librosa.istft(stft_denoised)
        
        return data_denoised
    
    def _apply_equalization(
        self, data: np.ndarray, sample_rate: int, bands: Dict[str, float]
    ) -> np.ndarray:
        """均衡器"""
        if len(data.shape) == 1:
            data = data.reshape(-1, 1)
        
        # 设计频段滤波器
        nyquist = sample_rate // 2
        
        # 低频滤波器 (60-250 Hz)
        low_freqs = [60, 250]
        b_low, a_low = signal.butter(2, [f/nyquist for f in low_freqs], btype='band')
        
        # 中低频滤波器 (250-500 Hz)
        low_mid_freqs = [250, 500]
        b_low_mid, a_low_mid = signal.butter(2, [f/nyquist for f in low_mid_freqs], btype='band')
        
        # 中频滤波器 (500-2000 Hz)
        mid_freqs = [500, 2000]
        b_mid, a_mid = signal.butter(2, [f/nyquist for f in mid_freqs], btype='band')
        
        # 中高频滤波器 (2000-4000 Hz)
        high_mid_freqs = [2000, 4000]
        b_high_mid, a_high_mid = signal.butter(2, [f/nyquist for f in high_mid_freqs], btype='band')
        
        # 高频滤波器 (4000-16000 Hz)
        high_freqs = [4000, min(16000, nyquist)]
        b_high, a_high = signal.butter(2, [f/nyquist for f in high_freqs], btype='band')
        
        # 应用滤波器和增益
        result = np.zeros_like(data)
        
        for ch in range(data.shape[1]):
            channel_data = data[:, ch]
            
            # 低频
            low_freq = signal.filtfilt(b_low, a_low, channel_data)
            result[:, ch] += low_freq * bands.get('low', 1.0)
            
            # 中低频
            low_mid_freq = signal.filtfilt(b_low_mid, a_low_mid, channel_data)
            result[:, ch] += low_mid_freq * bands.get('low_mid', 1.0)
            
            # 中频
            mid_freq = signal.filtfilt(b_mid, a_mid, channel_data)
            result[:, ch] += mid_freq * bands.get('mid', 1.0)
            
            # 中高频
            high_mid_freq = signal.filtfilt(b_high_mid, a_high_mid, channel_data)
            result[:, ch] += high_mid_freq * bands.get('high_mid', 1.0)
            
            # 高频
            high_freq = signal.filtfilt(b_high, a_high, channel_data)
            result[:, ch] += high_freq * bands.get('high', 1.0)
        
        return result.flatten() if data.shape[1] == 1 else result
    
    def _apply_dynamics_compression(
        self, data: np.ndarray, ratio: float, threshold_db: float
    ) -> np.ndarray:
        """动态压缩"""
        threshold_linear = np.power(10, threshold_db / 20)
        
        # 计算输入信号电平
        input_level = np.abs(data)
        
        # 应用压缩
        compressed_data = np.where(
            input_level > threshold_linear,
            threshold_linear + (input_level - threshold_linear) / ratio,
            input_level
        )
        
        # 保持原始相位
        compressed_data = np.sign(data) * compressed_data
        
        return compressed_data
    
    def _apply_stereo_enhancement(
        self, data: np.ndarray, width: float
    ) -> np.ndarray:
        """立体声增强"""
        if len(data.shape) == 1:
            # 单声道转立体声
            stereo_data = np.column_stack([data, data])
        else:
            stereo_data = data.copy()
        
        if stereo_data.shape[1] >= 2:
            left = stereo_data[:, 0]
            right = stereo_data[:, 1]
            
            # 计算中央和侧边信号
            mid = (left + right) / 2
            side = (left - right) / 2
            
            # 增强侧边信号
            side_enhanced = side * width
            
            # 重构左右声道
            left_enhanced = mid + side_enhanced
            right_enhanced = mid - side_enhanced
            
            # 限制幅度
            left_enhanced = np.clip(left_enhanced, -1.0, 1.0)
            right_enhanced = np.clip(right_enhanced, -1.0, 1.0)
            
            stereo_data[:, 0] = left_enhanced
            stereo_data[:, 1] = right_enhanced
        
        return stereo_data
    
    def _apply_pitch_correction(
        self, data: np.ndarray, sample_rate: int, pitch_shift: float
    ) -> np.ndarray:
        """音调校正"""
        if pitch_shift == 0:
            return data
        
        # 使用librosa进行音调调整
        shifted_data = librosa.effects.pitch_shift(
            data, sr=sample_rate, n_steps=pitch_shift
        )
        
        return shifted_data
    
    def _apply_tempo_adjustment(
        self, data: np.ndarray, sample_rate: int, tempo_factor: float
    ) -> np.ndarray:
        """节奏调整"""
        if tempo_factor == 1.0:
            return data
        
        # 使用librosa进行节奏调整
        adjusted_data = librosa.effects.time_stretch(data, rate=tempo_factor)
        
        return adjusted_data
    
    def _apply_reverberation(
        self, data: np.ndarray, sample_rate: int, room_size: float, damping: float
    ) -> np.ndarray:
        """混响效果"""
        # 简单的混响实现
        if room_size <= 0:
            return data
        
        # 创建混响 impulse response
        reverb_length = int(room_size * sample_rate)
        impulse_response = np.zeros(reverb_length)
        
        # 添加早期反射
        early_reflections = [0.01, 0.02, 0.03, 0.05]  # 秒
        for delay in early_reflections:
            delay_samples = int(delay * sample_rate)
            if delay_samples < reverb_length:
                impulse_response[delay_samples] = 0.3 * np.exp(-delay * damping)
        
        # 添加扩散混响
        for i in range(1, reverb_length):
            impulse_response[i] += 0.1 * np.exp(-i / (sample_rate * room_size) * damping)
            impulse_response[i] *= np.random.random() * 0.1
        
        # 应用混响
        reverberated_data = signal.fftconvolve(data, impulse_response, mode='same')
        
        # 混合原始信号和混响信号
        wet_level = min(room_size, 0.5)  # 混响强度
        dry_wet_mix = data * (1 - wet_level) + reverberated_data * wet_level
        
        return dry_wet_mix
    
    def _apply_clarity_enhancement(
        self, data: np.ndarray, sample_rate: int, strength: float
    ) -> np.ndarray:
        """清晰度增强"""
        # 使用高频增强提升清晰度
        
        # 设计高通滤波器
        nyquist = sample_rate // 2
        high_cutoff = 2000 / nyquist
        
        b_high, a_high = signal.butter(2, high_cutoff, btype='high')
        
        # 提取高频成分
        high_freq = signal.filtfilt(b_high, a_high, data)
        
        # 增强高频
        enhanced_high = high_freq * (1 + strength)
        
        # 与原始信号混合
        enhanced_data = data + enhanced_high * strength
        
        # 限制幅度
        enhanced_data = np.clip(enhanced_data, -1.0, 1.0)
        
        return enhanced_data
    
    def batch_enhance(
        self,
        audio_list: list,
        sample_rate: int,
        settings: EnhancementSettings
    ) -> list[EnhancementResult]:
        """批量音频增强"""
        results = []
        
        for i, audio_data in enumerate(audio_list):
            result = self.enhance_audio(audio_data, sample_rate, settings)
            result.metadata = result.metadata or {}
            result.metadata['batch_index'] = i
            results.append(result)
        
        return results
    
    def create_enhancement_preset(self, name: str, settings: EnhancementSettings):
        """创建增强预设"""
        # 这里可以实现预设管理
        pass
    
    def get_enhancement_preview(
        self,
        data: np.ndarray,
        sample_rate: int,
        enhancement_type: EnhancementType,
        strength: float = 0.5
    ) -> np.ndarray:
        """获取增强预览"""
        # 创建只包含指定增强类型的设置
        settings = EnhancementSettings()
        
        if enhancement_type == EnhancementType.NOISE_REDUCTION:
            settings.noise_reduction = True
            settings.noise_reduction_strength = strength
        elif enhancement_type == EnhancementType.EQUALIZATION:
            settings.equalization = True
            settings.eq_bands = {'mid': 1.5}  # 增强中频
        elif enhancement_type == EnhancementType.CLARITY:
            settings.clarity_enhancement = True
            settings.clarity_strength = strength
        # ... 其他类型的设置
        
        result = self.enhance_audio(data, sample_rate, settings)
        return result.enhanced_data if result.success else data
    
    def analyze_audio_quality(
        self, data: np.ndarray, sample_rate: int
    ) -> Dict[str, Any]:
        """分析音频质量"""
        try:
            analysis = {}
            
            # 计算基本统计信息
            analysis['rms'] = float(np.sqrt(np.mean(data ** 2)))
            analysis['peak'] = float(np.max(np.abs(data)))
            analysis['dynamic_range'] = float(20 * np.log10(analysis['peak'] / analysis['rms'])) if analysis['rms'] > 0 else 0
            
            # 频谱分析
            freqs, times, stft = signal.spectrogram(data, sample_rate)
            magnitude_spectrum = np.mean(np.abs(stft), axis=1)
            
            # 计算频谱质心
            spectral_centroid = np.sum(freqs * magnitude_spectrum) / np.sum(magnitude_spectrum)
            analysis['spectral_centroid'] = float(spectral_centroid)
            
            # 计算频谱带宽
            spectral_bandwidth = np.sqrt(
                np.sum(((freqs - spectral_centroid) ** 2) * magnitude_spectrum) / np.sum(magnitude_spectrum)
            )
            analysis['spectral_bandwidth'] = float(spectral_bandwidth)
            
            # 零交叉率
            zero_crossings = np.sum(np.diff(np.signbit(data)).astype(int))
            analysis['zero_crossing_rate'] = float(zero_crossings / len(data))
            
            # 建议的增强设置
            recommendations = []
            if analysis['dynamic_range'] > 30:
                recommendations.append('Consider dynamics compression')
            if analysis['spectral_centroid'] < 1000:
                recommendations.append('Consider high frequency enhancement')
            if analysis['zero_crossing_rate'] > 0.1:
                recommendations.append('Consider noise reduction')
            
            analysis['recommendations'] = recommendations
            
            return analysis
            
        except Exception as e:
            return {"error": f"Analysis failed: {str(e)}"}
